import numpy as np
import vecs_io
import bin_io
import os
import faiss


def delete_dir_if_exist(dire):
    if os.path.isdir(dire):
        # command = 'sudo rm -rf %s' % dire
        command = 'rm -rf %s' % dire
        print(command)
        os.system(command)


def groundtruth(base, query, k):
    base_dim = base.shape[1]
    index = faiss.IndexFlatL2(base_dim)
    index.add(base)
    gnd_distance, gnd_idx = index.search(query, k)
    print("search")
    return gnd_idx


if __name__ == '__main__':
    dataset_name = 'sift10K'

    old_dir = '/home/zhengbian/Dataset/siftsmall'

    basic_dir = '/home/zhengbian/nips-competition/data'
    dataset_dir = '%s/%s' % (basic_dir, dataset_name)
    delete_dir_if_exist(dataset_dir)
    os.mkdir(dataset_dir)

    base, dim = vecs_io.fvecs_read('%s/siftsmall_base.fvecs' % old_dir)
    base = base.astype(np.uint8)
    bin_io.bbin_write('%s/base.bbin' % dataset_dir, base)

    query, dim = vecs_io.fvecs_read('%s/siftsmall_query.fvecs' % old_dir)
    query = query.astype(np.uint8)
    bin_io.bbin_write('%s/query.bbin' % dataset_dir, query)

    gnd, dim = vecs_io.ivecs_read('%s/siftsmall_groundtruth.ivecs' % old_dir)
    gnd = gnd[:, :10].astype(np.uint32)
    bin_io.ibin_write('%s/gnd.ibin' % dataset_dir, gnd)

    print(base)
    print(query)
    print(gnd)
